#!/usr/bin/python
import os,re
try:
    import ConfigParser
except ImportError:     # Python 3.x
    import configparser as ConfigParser
    from functools import reduce
import json
import argparse

class JsonParser:
    def __init__(self, cf):
        self.json_file = cf

    def parse(self):
        group_base = dict()
        net_base = dict()
        fr = open(self.json_file, 'r')
        info = json.load(fr)
        for key, value in info.items():
            if(key != 'simulation'):
                continue
            sim_dict = value
        for key, value in sim_dict.items():
            if(key == 'Components'):
                comp_base = value
            elif(key == 'chiplets'):
                for k, v in value.items():
                    #print(k,v)
                    if(k == 'GroupNum' or k == 'Interconnect'):
                        net_base.update(v)
                group_base.update(self.group_basecfg(value, comp_base))

        #print(net_base)
        #print(group_base)
        return group_base, net_base

    #base configuration of group according to json
    def group_basecfg(self, group, comp):
        for key, value in group.items():
            if(key == 'GroupNum' or key == 'Interconnect'):
                continue
            if('CPU_type' not in value.keys()):
                print("*****Error: group %s has not specified CPU type"%key)
                sys.exit()
            cpu_type = value['CPU_type']
            value['CPU'] = comp[cpu_type]
            
            if('GPU_type' not in value.keys()):
                print("*****Error: group %s has not specified GPU type"%key)
                sys.exit()
            gpu_type = value['GPU_type']
            value['GPU'] = comp[gpu_type]

            if('Memory_type' not in value.keys()):
                print("*****Error: group %s has not specified MEM type"%key)
                sys.exit()
            mem_type = value['Memory_type']
            value['MEM'] = comp[mem_type]

        return group

    def get_np(self, output_file):
        all_gpu = 0
        group_base, net_base = self.parse()
        for key, value in net_base.items():
            if(len(re.findall(r'group\d+_type', key))):
                all_gpu += int(group_base[value]['GPU_count'])
        fw = open(output_file, 'w')
        if(all_gpu <= 1 ):
            np = 1
        else:
            np = all_gpu + int(net_base['num'])
        fw.write(str(np))
        os.putenv('np', str(all_gpu + int(net_base['num'])))
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-j", "--json", help = "specify json file", required = True)
    parser.add_argument("-o", "--output", help = "specify optput file", required = True)
    args = parser.parse_args()
    app = JsonParser(args.json) # TODO: change the json file name
    app.get_np(args.output)

